import os
import cv2
import numpy as np
import random
from tqdm import tqdm
import re
from glob import glob
def change_lightness(image):
    """Randomly change the lightness of an image."""
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    lightness_factor = random.uniform(0.6, 1.2)
    hsv[:, :, 2] = np.clip(hsv[:, :, 2] * lightness_factor, 0, 255)
    return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

def add_gaussian_noise(image):
    """Add Gaussian noise to an image."""
    prob = random.choice([0,0,1])
    if prob==0:
        return image
    row, col, ch = image.shape
    mean = 0
    var = random.uniform(0.01, 0.05)
    sigma = var ** 0.5
    gauss = np.random.normal(mean, sigma, (row, col, ch))
    noisy = image + gauss
    return np.clip(noisy, 0, 255).astype(np.uint8)

def add_gaussian_blur(image):
    """Add Gaussian blur to an image."""
    ksize = random.choice([ 1,1, 3,5])
    if ksize == 1:
        return image
    blurred = cv2.GaussianBlur(image, (ksize, ksize), 0)
    return blurred
def load_image(file_path):
    """Load a single image from the given path."""
    return cv2.imread(file_path, cv2.IMREAD_UNCHANGED)

def rotate_and_resize(image, mask, vis, angle, scale):
    """Rotate and resize an image, its mask, and its visualized mask."""
    h, w = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, scale)

    # Compute the bounding box for the rotated image
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
    new_w = int((h * sin) + (w * cos))
    new_h = int((h * cos) + (w * sin))

    # Adjust the rotation matrix to account for the new dimensions
    M[0, 2] += (new_w / 2) - center[0]
    M[1, 2] += (new_h / 2) - center[1]

    # Perform the rotation and resizing
    rotated_image = cv2.warpAffine(image, M, (new_w, new_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))
    rotated_mask = cv2.warpAffine(mask, M, (new_w, new_h), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
    rotated_vis = cv2.warpAffine(vis, M, (new_w, new_h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0))

    return rotated_image, rotated_mask, rotated_vis

def is_non_overlapping(mask, component_mask, x, y):
    """Check if a component can be placed at (x, y) without overlapping."""
    ch, cw = component_mask.shape[:2]
    if x + cw > mask.shape[1] or y + ch > mask.shape[0]:
        return False

    region = mask[y:y + ch, x:x + cw]
    region_non_zero = region > 0
    component_mask_non_zero = component_mask > 0
    overlap = np.sum(np.logical_and(region_non_zero, component_mask_non_zero))
    # overlap = np.sum(cv2.bitwise_and(region > 0, component_mask > 0))
    return overlap == 0

def match_component_luminance_lab(component_rgb, background_rgb, mask,meaningful_mask):
    """
    将元件图像在LAB空间下的亮度（L通道）调整为与背景相似，以提升融合自然度。

    参数:
        component_rgb: 元件图 (H, W, 3), RGB或BGR格式均可
        background_rgb: 背景图 (H, W, 3), RGB或BGR格式均可
        mask: uint8灰度图 (H, W), 非0区域为元件

    返回:
        调整后的元件图像 (H, W, 3)
    """
    # 转为 LAB 空间
    component_lab = cv2.cvtColor(component_rgb, cv2.COLOR_BGR2LAB)
    background_lab = cv2.cvtColor(background_rgb, cv2.COLOR_BGR2LAB)

    # 分离通道
    c_L, c_A, c_B = cv2.split(component_lab)
    b_L, _, _ = cv2.split(background_lab)

    # 仅在mask区域统计亮度分布
    component_mask = mask > 0
    meaningful_mask = meaningful_mask > 0
    combined_mask = component_mask & meaningful_mask
    if np.sum(combined_mask) == 0:
        return component_rgb  # 避免空mask时除零

    c_mean = np.mean(c_L[component_mask])
    c_std = np.std(c_L[component_mask]) + 1e-5  # 避免除0
    b_mean = np.mean(b_L[combined_mask])  # 背景中对应区域亮度
    b_std = np.std(b_L[combined_mask]) + 1e-5

    # 匹配亮度：线性变换
    matched_L = (c_L.astype(np.float32) - c_mean) * (b_std / c_std) + b_mean
    matched_L = np.clip(matched_L, 0, 255).astype(np.uint8)

    # 合并回LAB并转回BGR
    matched_lab = cv2.merge([matched_L, c_A, c_B])
    matched_bgr = cv2.cvtColor(matched_lab, cv2.COLOR_LAB2BGR)

    return matched_bgr
def place_component(board, board_mask, board_cpimask, board_vis_mask,
                    component, component_mask, cpmask, component_vis_mask,
                    x, y,feather=False,diag=False,diag_param=[]):

    """Place a component on the board and update mask and visualized mask."""
    ch, cw = component.shape[:2]
    foreground_indices = component_mask > 0
    # 调整元件亮度和对比度
    if diag:
         # 加载有意义区域的mask
        meaningful_mask_path = os.path.join("data_zoo/pcbbgpool/train/refine", os.path.basename(diag_param)+'_effect_mask.png')
        meaningful_mask = cv2.imread(meaningful_mask_path, cv2.IMREAD_GRAYSCALE)
        meaningful_maskcoords = cv2.findNonZero(meaningful_mask)
        maskbbox= cv2.boundingRect(meaningful_maskcoords)
        meaningful_mask=meaningful_mask[maskbbox[1]:maskbbox[1]+maskbbox[3],maskbbox[0]:maskbbox[0]+maskbbox[2]]
        meaningful_mask = meaningful_mask[y:y + ch, x:x + cw]
        component = match_component_luminance_lab(component,board[y:y + ch, x:x + cw], component_mask,meaningful_mask)
     
     # Feather technique: Create a feathered mask
    if feather:
        feathered_mask = cv2.GaussianBlur(foreground_indices.astype(np.float32), (31, 31), 0)
        feathered_mask = feathered_mask / np.max(feathered_mask)
        for c in range(3):
            board[y:y + ch, x:x + cw, c][foreground_indices]=(
                feathered_mask[foreground_indices] * component[:, :, c][foreground_indices] +
                (1 - feathered_mask[foreground_indices]) * board[y:y + ch, x:x + cw, c][foreground_indices]
            ).astype(np.uint8)
        fb=board
    else:
        for c in range(3):
            board[y:y + ch, x:x + cw, c][foreground_indices] = component[:, :, c][foreground_indices]

    board_mask[y:y + ch, x:x + cw][foreground_indices] = np.maximum(
        board_mask[y:y + ch, x:x + cw][foreground_indices], component_mask[foreground_indices]
    )
    board_cpimask[y:y + ch, x:x + cw][foreground_indices] = np.maximum(
        board_cpimask[y:y + ch, x:x + cw][foreground_indices], cpmask[foreground_indices]
    )

    board_vis_mask[y:y + ch, x:x + cw][foreground_indices] = component_vis_mask[foreground_indices]

def synthesize_pcb_dataset(empty_board_path, component_root, output_folder, num_images, min_scale=0.5, max_scale=1.5,feather=True):
    """Generate synthetic PCB dataset."""
    def get_bbox_from_mask(mask, bgwh, shift):
        """
        从掩码图像中获取掩码部分的边界框。

        :param mask: 输入的掩码图像，形状为 (h, w)
        :return: 边界框信息 (x, y, width, height) xywh
        """
        # 找到掩码中非零像素的坐标
        rows, cols = np.nonzero(mask)
        bw, bh = bgwh
        sx, sy = shift

        if len(rows) == 0 or len(cols) == 0:
            return (0, 0, 0, 0)

        # 计算边界框的左上角坐标和宽高
        x = np.min(cols)
        y = np.min(rows)
        width = np.max(cols) - x + 1
        height = np.max(rows) - y + 1

        return ((x + width / 2 + sx) / bw, (y + height / 2 + sy) / bh, width / bw, height / bh)

    def get_component_type(component_path):
        type_name = component_path[len(component_root):]
        type_name = os.path.dirname(type_name)
        pattern = r'JRS(\d+)[a-zA-Z].*'
        match = re.match(pattern, type_name)
        number = match.group(1)
        return int(number)

    # Load the empty board
    pcb_bg_paths = glob(os.path.join(empty_board_path, "*.jpg"))

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    component_paths = []  # len=1948
    for root, dirs, files in os.walk(component_root):
        for file in files:
            if file.endswith(('.png', '.jpg', '.jpeg')) and 'label' not in root:
                component_paths.append(os.path.join(root, file))

    image_files = glob(os.path.join(output_folder, "image", "*.png"))
    start_index = len(image_files)
    for i in tqdm(range(start_index,num_images)):
        pcb_bg_path = random.choice(pcb_bg_paths)
        board = cv2.imread(pcb_bg_path)
        board=change_lightness(board)
        board = add_gaussian_blur(board)
        board = add_gaussian_noise(board)
        
        bh,bw,_=board.shape
        # max_n_cp = int(bh * bw / 300 / 400 * 1.2)
        # min_n_cp = int(bh * bw / 300 / 400 * 0.5)
        max_n_cp = int(bh * bw / 300 / 400 * 0.8)
        min_n_cp = int(bh * bw / 300 / 400 * 0.4)
        board_mask = np.zeros(board.shape[:2], dtype=np.uint8)
        board_cpimask = board_mask.copy()
        vis_mask = np.zeros_like(board)
        yolo_labels = []

        used_components = []
        num_components = random.randint(min_n_cp, max_n_cp)  # Random number of components per board
        error_guest = [
            "JRS183R0603/_a_SaveData__20241013161630397_Save Data [14837, 35647]",
            "JRS156R0603/_a_SaveData__20241013161630396_Save Data [22823, 35559]",
            "JRS157R0603/_a_SaveData__20241013161630397_Save Data [18766, 35613]",
            "JRS116PGLQFP/_a_SaveData__20241013155754673_SaveData [23447, 102811]",
            "JRS104TSOP50",
            "JRS116PGLQFP",
            "JRS82R1812/_a_SaveData__20241013155402990_Save Data [109592, 18847]",
        ]
        # if feather:
        #     # Save feathered component
        #     feathered_board = board.copy()
        #     feathered_board_mask = board_mask.copy()
        #     feathered_board_cpimask = board_cpimask.copy()
        #     feathered_vis_mask = vis_mask.copy()
        for it in range(num_components):
            component_path = random.choice(component_paths)
            label_path = component_path.replace('/_a_SaveData', '/label/_a_SaveData').replace('jpg', 'png')
            vis_path = component_path.replace('/_a_SaveData', '/label/_a_SaveData').replace('.jpg', '_pseudo.png')

            log_file_path = 'og.log'
            error_flag = False
            for error_path in error_guest:
                if error_path in component_path:
                    it -= 1
                    error_flag = True
                    break
            if error_flag:
                continue

            component = load_image(component_path)
            component = add_gaussian_blur(component)
            component = add_gaussian_noise(component)
            
            component_mask = load_image(label_path)
            component_vis_mask = load_image(vis_path)

            if component is None or component_mask is None:
                warning_message = f"Warning: Failed to load data for {component_path}"
                status_message = f"{component is None}, {component_mask is None}"

                with open(log_file_path, 'a') as log_file:
                    log_file.write(warning_message + '\n')
                    log_file.write(status_message + '\n')
                continue

            angle = random.uniform(0, 360)
            scale = random.uniform(min_scale, max_scale)

            component, component_mask, component_vis_mask = rotate_and_resize(
                component, component_mask, component_vis_mask, angle, scale / 2
            )
            

            ch, cw = component.shape[:2]
            do_it_more = max(1, 1024 * 1024 / ch / cw - 1)
            for _ in range(100):  # Maximum attempts to place the component
                x = random.randint(0, board.shape[1] - cw)
                y = random.randint(0, board.shape[0] - ch)

                if is_non_overlapping(board_mask, component_mask, x, y):
                    bbox_xywh = get_bbox_from_mask(component_mask, (bw,bh), (x, y))
                    class_id = get_component_type(component_path)

                    cpmask = np.copy(component_mask)
                    non_zero_indices = cpmask != 0
                    cpmask[non_zero_indices] = class_id

                    place_component(
                        board, board_mask, board_cpimask, vis_mask,
                        component, component_mask, cpmask, component_vis_mask,
                        x, y,feather=True,diag=True,diag_param=pcb_bg_path
                    )
                    # place_component(
                    #     feathered_board, feathered_board_mask, feathered_board_cpimask, feathered_vis_mask,
                    #     component, component_mask, cpmask, component_vis_mask,
                    #     x, y,feather=True
                    # )
                    # place_component(
                    #     feathered_board, feathered_board_mask, feathered_board_cpimask, feathered_vis_mask,
                    #     component, component_mask, cpmask, component_vis_mask,
                    #     x, y,diag=True,diag_param=pcb_bg_path
                    # )
                    used_components.append(component_path)
                    yolo_labels.append((class_id,) + bbox_xywh)
                    break

            del component, component_mask, component_vis_mask

        # Save outputs
        dir_list = ["image", "mask", "vis", "cplist", "cpbbox", "cpimask"]
        for dir_name in dir_list:
            if not os.path.exists(os.path.join(output_folder, dir_name)):
                os.makedirs(os.path.join(output_folder, dir_name))
        output_image_path = os.path.join(os.path.join(output_folder, "image"), f"pcb_{i:04d}.png")
        output_mask_path = os.path.join(os.path.join(output_folder, "mask"), f"pcb_{i:04d}_mask.png")
        output_cpimask_path = os.path.join(os.path.join(output_folder, "cpimask"), f"pcb_{i:04d}_cpimask.png")
        output_vis_mask_path = os.path.join(os.path.join(output_folder, "vis"), f"pcb_{i:04d}_vis_mask.png")
        output_components_path = os.path.join(os.path.join(output_folder, "cplist"), f"pcb_{i:04d}_components.txt")
        output_bbox_path = os.path.join(os.path.join(output_folder, "cpbbox"), f"pcb_{i:04d}_cpbbox.txt")
        cv2.imwrite(output_image_path, board)
        cv2.imwrite(output_mask_path, board_mask)
        cv2.imwrite(output_cpimask_path, board_cpimask)
        cv2.imwrite(output_vis_mask_path, vis_mask)
        
        # Save outputs for feathered board
        output_feathered_image_path = os.path.join(os.path.join(output_folder, "image"), f"pcb_{i:04d}_feathered.png")
        # output_feathered_mask_path = os.path.join(os.path.join(output_folder, "mask"), f"pcb_{i:04d}_feathered_mask.png")
        # output_feathered_cpimask_path = os.path.join(os.path.join(output_folder, "cpimask"), f"pcb_{i:04d}_feathered_cpimask.png")
        # output_feathered_vis_mask_path = os.path.join(os.path.join(output_folder, "vis"), f"pcb_{i:04d}_feathered_vis_mask.png")
        # cv2.imwrite(output_feathered_image_path, feathered_board)
        # cv2.imwrite(output_feathered_mask_path, feathered_board_mask)
        # cv2.imwrite(output_feathered_cpimask_path, feathered_board_cpimask)
        # cv2.imwrite(output_feathered_vis_mask_path, feathered_vis_mask)


        with open(output_bbox_path, 'w') as f2:
            for label in yolo_labels:
                bbox_label = label[:5]
                label_str = " ".join(str(i) for i in bbox_label)
                f2.write(label_str + "\n")

        with open(output_components_path, 'w') as f:
            f.write('\n'.join(used_components))

# Example usage
empty_board_path = "data_zoo/jrspcbbg"
component_root = "/home/lhx/pcbsyn/data_zoo/PCBA/20241209_PCBA样板标图数据-华科用/"
output_folder = "synpcbwithsimplejrsbg"
synthesize_pcb_dataset(
    empty_board_path=empty_board_path,
    component_root=component_root,
    output_folder=output_folder,
    num_images=60,
    min_scale=1.1,
    max_scale=0.9
)